import pymysql
import sshtunnel
from config import isLocalDB

def connect_to_database():
    if isLocalDB:
        ssh_tunnel = None
        db = pymysql.connect(
            host="localhost",
            user="ucas_user",
            password="ucas_password",
            database="ucas_food",
            charset="utf8",
        )
        print("Database connection established successfully.")
    else:
        # SSH 隧道配置
        ssh_tunnel = sshtunnel.SSHTunnelForwarder(
            ("124.70.73.194", 22),  # SSH 服务器地址和端口
            ssh_username="root",  # SSH 用户名
            ssh_password="!@#123qwe",  # SSH 密码
            remote_bind_address=("localhost", 3306),  # 远程数据库地址和端口
        )
        try:
            # 启动 SSH 隧道
            ssh_tunnel.start()
            print("SSH tunnel started successfully.")

            # 连接到数据库
            db = pymysql.connect(
                host="127.0.0.1",  # 本地地址
                port=ssh_tunnel.local_bind_port,  # 本地端口
                user="ucas_food",  # 数据库用户名
                password="123456",  # 数据库密码
                database="ucas_food",  # 数据库名称
                charset="utf8",
            )
            print("Database connection established successfully.")

            return db, ssh_tunnel
        except Exception as e:
            print(f"Error connecting to database: {e}")
            raise
    return db, ssh_tunnel


def close_database_connection(db, ssh_tunnel):
    db.close()
    if ssh_tunnel is not None:
        ssh_tunnel.stop()
